
import pickle as pickle
import numpy as np
import matplotlib.pyplot as plt
from matplotlib.colors import BoundaryNorm, ListedColormap

infile = open('results-eddy1e5-shield300-krON-jan25.txt','r')
temp = infile.readline().split('Result:')
a = []
while len(temp)>1:
    a += [float(temp[1])]
    temp = infile.readline().split('Result:')
a = np.array(a)
a = np.reshape(a,(15,20))

with open('znodust.pickle','rb') as handle:
    dat = pickle.load(handle)
    
# blank out the rainy cases
import pickle as pickle
with open('moistheight20.pickle','rb') as handle:
    Z = pickle.load(handle)
x,y=np.shape(Z)
for ix in range(0,x):
    for iy in range(0,y):
        if Z[ix,iy] == 3:
            a[ix,iy] = 0.0

# plot albedo
plt.clf()
colors = ['k','white','navy','mediumblue','royalblue','cornflowerblue','mediumturquoise','khaki','peachpuff','lightsalmon','salmon','indianred','red']
cmap = ListedColormap(colors)
boundaries = [-0.01,0.1,0.5,0.55,0.6,0.65,0.7,0.75,0.8,0.85,0.9,0.95]
norm = BoundaryNorm(boundaries, ncolors=len(colors), extend='both')
print(dat.keys())
plt.figure(figsize=(8, 6))
contour = plt.contourf(dat['fv'], dat['ts'], a, levels=boundaries, cmap=cmap, norm=norm)
plt.colorbar(contour, ticks=boundaries[1:])
plt.gca().invert_yaxis()
plt.xscale('log')
plt.xlabel("Maximum Mixing Ratio of Water Vapor")
plt.ylabel("Surface Temperature")
plt.savefig('a.png',dpi=300)

######################################################################################
# works to here

# plot asr v olr
Ts_val_a = np.arange(300,1001,50,dtype='float')
nT = np.size(Ts_val_a)
nH = 20 # 14    
fv_val_a = np.array([1.000e-06,1.000e-05,1.000e-04,2.258e-04,5.097e-04,1.151e-03,2.597e-03,5.864e-03,1.324e-02,2.989e-02,6.747e-02,8.000e-02, \
                     9.500e-02, 1.100e-01, 1.300e-01, 1.700e-01, 2.000e-01,\
                     1.523e-01,3.438e-01,7.762e-01])

T = np.zeros((nT,nH,110))
p = np.zeros((nT,nH,110))
fH2O = np.zeros((nT,nH,110))
fH2O_surf = np.zeros((nT,nH))
fCO2 = np.zeros((nT,nH,110))

sigma = 5.67e-8
with open('tpmore.pickle','rb') as handle:
    dat_Tp = pickle.load(handle)

for i in range(0,nT):
    Ts_val = Ts_val_a[i]
    for j in range(0,nH):
        fv_val = fv_val_a[j]

        fv_str = "{:.3e}".format(fv_val)
        Ts_str = "{:.1f}".format(Ts_val)
        fname = 'exovenus-fv'+fv_str+'-Ts'+Ts_str+'.out-1'
        try:
            one_prof = dat_Tp[fname]
            T[i,j,:] = one_prof['T']
            p[i,j,:] = one_prof['P']*1e-3 # convert mBar to bar
            fH2O[i,j,:] = one_prof['H2O']/one_prof['den']
            fCO2[i,j,:] = one_prof['CO2']/one_prof['den']
        except:
            print(fname)


fH2O_surf[:,:] = fH2O[:,:,-1]

kapH2O = 0.01 #
kapCO2 = 0.001 
costh = 0.5
g = 8.87 # 
p_em_H2O = (g*costh/kapH2O)/1e5 
p_em_CO2 = (g*costh/kapCO2)/1e5

T_em_H2O = np.zeros((nT,nH))
T_em_CO2 = np.zeros((nT,nH))

T_em = np.zeros((nT,nH))
p_em = np.zeros((nT,nH))
fH2O_em = np.zeros((nT,nH))
OLR = np.zeros((nT,nH))

for i in range(0,nT):
    Ts_val = Ts_val_a[i]
    for j in range(0,nH):
        fv_val = fv_val_a[j]

        p_H2O_int = fH2O[i,j,:]*p[i,j,:]
        p_CO2_int = fCO2[i,j,:]*p[i,j,:]
        T_int    = T[i,j,:]

        fH2O_int = fH2O[i,j,:]
        T_em_H2O[i,j] = np.interp(p_em_H2O, p_H2O_int[::-1], T_int[::-1]) 
        T_em_CO2[i,j] = np.interp(p_em_CO2, p_CO2_int[::-1], T_int[::-1]) 
        T_em[i,j] = np.min([T_em_H2O[i,j],T_em_CO2[i,j]])
        p_em[i,j] = p_em_H2O*(T_em_H2O[i,j]<T_em_CO2[i,j]) + p_em_CO2*(T_em_H2O[i,j]>T_em_CO2[i,j])

        OLR[i,j]  = sigma*T_em[i,j]**4


F0 = 0.75*1362/0.723**2
ASR = (1-a)*F0/4


import pickle as pickle
with open('moistheight20.pickle','rb') as handle:
    Z = pickle.load(handle)
x,y=np.shape(Z)
for ix in range(0,x):
    for iy in range(0,y):
        if Z[ix,iy] == 3:
            ASR[ix,iy] = -9000

colors = ['white','k','navy','mediumblue','royalblue','cornflowerblue','mediumturquoise','khaki','peachpuff','lightsalmon','salmon','indianred','red']
boundaries = [-10000,-200,-40,-20,-10,-5,0,5,10,20,40,250]
norm = BoundaryNorm(boundaries, ncolors=len(colors), extend='both')
plt.figure(figsize=(8, 6))
contour = plt.contourf(dat['fv'], dat['ts'], ASR-OLR, levels=boundaries, cmap=cmap, norm=norm)
plt.colorbar(contour, ticks=boundaries[1:])
plt.gca().invert_yaxis()
plt.xscale('log')
plt.xlabel("Maximum Mixing Ratio of Water Vapor")
plt.ylabel("Surface Temperature")
plt.savefig('df10.png',dpi=300)
plt.clf()


###################################
# cartoon here

plt.clf()
with open('zcartoon.pickle','rb') as handle:
    z = pickle.load(handle)

plt.figure(figsize=(8,6))
boundaries = [0.5,1.5,2.5,3.5]
colors = ['white', 'coral', 'gold', 'cornflowerblue','k']
cmap = ListedColormap(colors)
norm = BoundaryNorm(boundaries, ncolors=len(colors), extend='both')
#contour = plt.contourf(dat['fv'], dat['ts'], z, levels=boundaries, cmap=cmap, norm=norm)                                                                                                  
mesh = plt.pcolormesh(dat['fv'], dat['ts'], z, shading='nearest', cmap=cmap, norm=norm)
plt.colorbar(contour, ticks=boundaries[1:])
plt.gca().invert_yaxis()
plt.contour(dat['fv'], dat['ts'], ASR-OLR, levels=[0], colors='white', linewidths=2,linestyles='-')
plt.xscale('log')
plt.xlabel("Maximum Mixing Ratio of Water Vapor")
plt.ylabel("Surface Temperature")
plt.savefig('cartoonjan6.png',dpi=300)

#####################################

# tau h2o

with open('tauh2.pickle','rb') as handle:
    tau = pickle.load(handle)


x,y=np.shape(Z)
for ix in range(0,x):
    for iy in range(0,y):
        if Z[ix,iy] == 3:
            tau[ix,iy] = 0.01
            
plt.clf()
colors = ['k','white','white','navy','mediumblue','royalblue','cornflowerblue','mediumturquoise','khaki','peachpuff','lightsalmon','salmon','indianred','red']
cmap = ListedColormap(colors)
boundaries = [-10,-1,10,3e1,5e1,1e2,3e2,1e3,3e3,1e4,3e4,1e5,3e5]
norm = BoundaryNorm(boundaries, ncolors=len(colors), extend='both')
print(dat.keys())

tau[-3,-1] = tau[-2,-1]
plt.figure(figsize=(8, 6))
contour = plt.contourf(dat['fv'], dat['ts'],tau, levels=boundaries, cmap=cmap, norm=norm)

print(tau[-3,-1])
tau[-3,-1] = tau[-2,-1]
print(tau)
plt.colorbar(contour, ticks=boundaries[1:])
plt.gca().invert_yaxis()
plt.xscale('log')
plt.xlabel("Maximum Mixing Ratio of Water Vapor")
plt.ylabel("Surface Temperature")
plt.savefig('tauh2.png',dpi=300)
